# -*- coding: utf-8 -*-

# @Project : fastapi-tutorial
# @Date    : 20240322-0139
# @Author  : robin

import os
import sys
from dotenv import load_dotenv
# Load environment variables from openai.env file
load_dotenv("../.env")


from langchain.chains import load_chain

# chain = load_chain("D:/project/samer/workspace/fastapi-tutorial/src/hub/chain.json")
# print(chain.run("2+6等于几?"))

# chain = load_chain("D:/project/samer/workspace/fastapi-tutorial/src/hub/hello_chain.json")
# print(chain.run("男人"))

from typing import List, Dict, Any, Optional
from langchain.callbacks.manager import (
    CallbackManagerForChainRun
)
from langchain.chains.base import  Chain
from langchain.prompts.base import BasePromptTemplate
from langchain.base_language import  BaseLanguageModel

# CustomChain
# 当通用链不满足的时候，可以自行构建来实现特定的目的
from typing import List, Dict, Any, Optional
from langchain.callbacks.manager import (
    CallbackManagerForChainRun
)
from langchain.chains.base import  Chain
from langchain.prompts.base import BasePromptTemplate
from langchain.base_language import  BaseLanguageModel


class wiki_article_chain(Chain):
    """开发一个wiki文章生成器"""
    prompt: BasePromptTemplate
    llm: BaseLanguageModel
    out_key: str = "text"

    @property
    def input_keys(self) -> List[str]:
        """将返回Prompt所需的所有键"""
        return self.prompt.input_variables

    @property
    def output_keys(self) -> List[str]:
        """将始终返回text键"""
        return [self.out_key]

    def _call(
            self,
            inputs: Dict[str, Any],
            run_manager: Optional[CallbackManagerForChainRun] = None,
    ) -> Dict[str, Any]:
        """运行链"""
        prompt_value = self.prompt.format_prompt(**inputs)
        # print("prompt_value:",prompt_value)
        response = self.llm.generate_prompt(
            [prompt_value], callbacks=run_manager.get_child() if run_manager else None
        )
        # print("response:",response)
        if run_manager:
            run_manager.on_text("文章写完得话") # wiki article is written
        return {self.out_key: response.generations[0][0].text}

    @property
    def _chain_type(self) -> str:
        """链类型"""
        return "wiki_article_chain"

from langchain.chat_models import  ChatOpenAI
from langchain.prompts import  PromptTemplate
from langchain.llms import OpenAI

chain = wiki_article_chain(
    prompt=PromptTemplate(
        template= "关于这个{topic}主题,写一个好玩得笑话", #"写一篇关于{topic}的维基百科形式的文章",
        input_variables=["topic"]
    ),
    # llm=OpenAI(
    #     temperature=0,
    #     model_name = "gpt-3.5-turbo-instruct"
    # ),
    llm=ChatOpenAI(
        temperature=0,
        model_name = "gpt-3.5-turbo"
    ),
)

result = chain.run({"topic":"男人"})
print(result)

